Skip to content

[atom-vllm] adapt DeepSeek V4 MTP for vLLM plugin - #1372

Merged
zejunchen-zejun merged 9 commits into
mainfrom
hexwang/ds_v4_mtp
Jul 6, 2026
Merged

[atom-vllm] adapt DeepSeek V4 MTP for vLLM plugin#1372
zejunchen-zejun merged 9 commits into
mainfrom
hexwang/ds_v4_mtp

Conversation

@whx-sjtu

@whx-sjtu whx-sjtu commented Jun 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

  • Wire DeepSeek V4 MTP into the ATOM vLLM plugin, including draft model registration, target/draft sharing, V4 proxy cache handling, and MTP-specific metadata adaptation.
  • Full Graph support for MTP.

Test Result

DeepSeek V4 Flash, TP4

Full gsm8k (1319 problems, 20-shot, concurrency 64):

  |Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
  |gsm8k|      3|flexible-extract|    20|exact_match|↑  |0.9416|±  |0.0065|
  |     |       |strict-match    |    20|exact_match|↑  |0.9416|±  |0.0065|

MTP=3 acceptance from the gsm8k run (48,367 draft steps × 3 tokens = 145,101 draft tokens; 77,360 accepted):

  ┌─────────────────────────────────────────────────┬──────────────────────┐
  │ Metric                                          │ Value                │
  ├─────────────────────────────────────────────────┼──────────────────────┤
  │ Overall acceptance rate (accepted/draft tokens) │ 53.3% (77360/145101) │
  ├─────────────────────────────────────────────────┼──────────────────────┤
  │ Mean accepted draft tokens per step             │ 1.60 (77360/48367)   │
  ├─────────────────────────────────────────────────┼──────────────────────┤
  │ Mean acceptance length (incl. bonus token)      │ ≈2.60 tokens/step    │
  └─────────────────────────────────────────────────┴──────────────────────┘

Per-position acceptance (how often the k-th speculated token is accepted):


  ┌──────────┬─────────────────────┐
  │ Position │ Acceptance          │
  ├──────────┼─────────────────────┤
  │ 0        │ 90.0% (43553/48367) │
  ├──────────┼─────────────────────┤
  │ 1        │ 52.2% (25249/48367) │
  ├──────────┼─────────────────────┤
  │ 2        │ 17.7% (8558/48367)  │
  └──────────┴─────────────────────┘

@whx-sjtu whx-sjtu changed the title adapt DeepSeek V4 MTP for vLLM plugin [atom-vllm] adapt DeepSeek V4 MTP for vLLM plugin Jun 26, 2026
@whx-sjtu
whx-sjtu marked this pull request as draft June 26, 2026 10:59
@whx-sjtu
whx-sjtu force-pushed the hexwang/ds_v4_mtp branch from ed000a4 to 987cc28 Compare June 26, 2026 11:04
@zufayu
zufayu requested a review from yitingw1 June 30, 2026 02:30
whx-sjtu added 4 commits June 30, 2026 08:43
@whx-sjtu
adapt v4 mtp
fd25521 
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
@whx-sjtu
FULL Graph support
91361de 
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
@whx-sjtu
fix graph
8b7de60 
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
@whx-sjtu
fix memory access & accpetance drop
5fa130c 
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
@whx-sjtu
whx-sjtu force-pushed the hexwang/ds_v4_mtp branch from e8b6483 to 5fa130c Compare July 1, 2026 02:59
@whx-sjtu
refactor
5c99e50 
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
@whx-sjtu
whx-sjtu marked this pull request as ready for review July 1, 2026 04:27
@whx-sjtu
remove redundant envs
d0dff32 
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
@whx-sjtu
whx-sjtu requested a review from zejunchen-zejun July 1, 2026 09:33
zejunchen-zejun
zejunchen-zejun reviewed Jul 1, 2026
View reviewed changes
Comment thread atom/plugin/vllm/register.py Outdated
Comment thread atom/plugin/vllm/deepseek_v4_bridge.py Outdated
if getattr(common_attn_metadata, "max_query_len", 0) == 1:
max_q = int(getattr(common_attn_metadata, "max_query_len", 0) or 0)
decode_q = 1 + max(0, int(num_spec_tokens))
if 1 <= max_q <= decode_q:

@zejunchen-zejun zejunchen-zejun Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if there is a short prefill request or chunked prefill request in the batch, whose max_q is very small, it seems this request could be wrongly treated as DECODE

@whx-sjtu whx-sjtu Jul 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good finding. I will add extra check here to filter out prefill mixed batches.

whx-sjtu added 2 commits July 1, 2026 13:16
@whx-sjtu
update register.py
069c9e4 
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
@whx-sjtu
extra check for decode
5b7f140 
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
@whx-sjtu
whx-sjtu force-pushed the hexwang/ds_v4_mtp branch from 61ff620 to 5b7f140 Compare July 2, 2026 09:15
yitingw1
yitingw1 approved these changes Jul 3, 2026
View reviewed changes

@yitingw1 yitingw1 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@zejunchen-zejun
zejunchen-zejun merged commit 6ba110b into main Jul 6, 2026
10 of 12 checks passed
@zejunchen-zejun
zejunchen-zejun deleted the hexwang/ds_v4_mtp branch July 6, 2026 05:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@valarLip valarLip valarLip approved these changes
@yitingw1 yitingw1 yitingw1 approved these changes
@zejunchen-zejun zejunchen-zejun zejunchen-zejun approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@whx-sjtu @valarLip @yitingw1 @zejunchen-zejun